import torch.nn as nn

from layers import GCN, AvgReadout, Discriminator, ClusterNet, DiscriminatorCluster


class GIC(nn.Module):
    def __init__(self, nodes, n_input, hidden, activation, num_clusters, beta):
        super(GIC, self).__init__()
        self.gcn = GCN(n_input, hidden, activation)
        self.read = AvgReadout()
        self.sigm = nn.Sigmoid()
        self.disc = Discriminator(hidden)
        self.disc_c = DiscriminatorCluster(hidden, hidden, nodes, num_clusters)
        self.beta = beta
        self.cluster = ClusterNet(hidden, num_clusters)

    def forward(self, pos, neg, adj, sparse, msk, samp_bias1, samp_bias2, cluster_tmp):
        h_1 = self.gcn(pos, adj, sparse)

        self.beta = cluster_tmp
        z, s = self.cluster(h_1[-1, :, :], cluster_tmp)
        z_t = s @ z
        c_2 = z_t
        c_2 = self.sigm(c_2)

        c = self.read(h_1, msk)
        c = self.sigm(c)
        c_x = c.unsqueeze(1)
        c_x = c_x.expand_as(h_1)

        h_2 = self.gcn(neg, adj, sparse)

        ret1 = self.disc(c_x, h_1, h_2, samp_bias1, samp_bias2)
        ret2 = self.disc_c(c_2, h_1[-1, :, :], h_1[-1, :, :], h_2[-1, :, :], samp_bias1, samp_bias2)
        return ret1, ret2

    def embed(self, seq, adj, sparse, msk):
        h_1 = self.gcn(seq, adj, sparse)
        c = self.read(h_1, msk)

        z, s = self.cluster(h_1[-1, :, :], self.beta)
        h = s @ z

        return h_1.detach(), h.detach(), c.detach(), z.detach()
